//
// Created by lanlu on 2025/8/11.
//

#include <stack>
#include <unordered_map>
using namespace std;

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;

    TreeNode() : val(0), left(nullptr), right(nullptr){}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr){}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right){}
};

class Solution {
public:
    // 使用一个容器 存储节点以及高度
    unordered_map<TreeNode*, int> height;

    TreeNode* balanceBST(TreeNode* root) {
        // 一颗新的平衡二叉树
        TreeNode * newRoot = nullptr;

        // 非递归的先序遍历
        stack<TreeNode*> st;
        TreeNode * cur = root;
        while (cur != nullptr || !st.empty()) {
            if (cur != nullptr)
            {
                st.push(cur);
                // 将此节点插入到平衡二叉树中
                newRoot = insert(newRoot, cur -> val);
                cur = cur->left;
            }
            else
            {
                cur = st.top();
                st.pop();
                cur = cur -> right;
            }
        }

        return newRoot;
    }

    // 插入前 传入当前的根节点 和 要插入的节点值
    // 插入后返回新的根节点
    TreeNode * insert(TreeNode * root, int val)
    {
        // 如果本身为空 直接插入一个节点
        if (root == nullptr)
        {
            root = new TreeNode(val);
            // 节点的高度需要更新
            height[root] = 1;
            return root;
        }

        // 如果根节点不为空 比较新插入的值 和 根节点的值相比 如果更大 那么在右子树 如果更小 那么在左子树
        TreeNode * node = root;
        int cmp = val - node -> val;
        if (cmp < 0)
        {
            // 左子树
            // 如果要插入到左子树中 就是递归调用当前函数 并且将处理结果赋给左子树
            node -> left = insert(node -> left, val);
            // 如果左右子树的高度差超过1 需要旋转
            // 如果插入到左子树 左子树可能变高 所以不需要求绝对值
            if (getHeight(node -> left) - getHeight(node -> right) > 1)
            {
                // 如果插入的节点在左子树的右边，需要先左旋再右旋
                // 要注意 左旋是针对左孩子旋转  右旋是针对自身
                if (val > node -> left -> val)
                {
                    // 左旋 针对左孩子 所以左孩子的根节点变化
                    node -> left = leftRotate(node -> left);
                }
                // 如果插入的节点在左子树的左边，直接右旋
                // 针对自身 所以自身的根节点变化
                node = rightRotate(node);
            }
        }
        else if (cmp > 0)
        {
            // 右子树
            node -> right = insert(node -> right, val);
            if (getHeight(node -> right) - getHeight(node -> left) > 1)
            {
                // 如果插入的节点在右子树的左边，需要先右旋再左旋
                // 要注意 右旋是针对右孩子旋转  左旋是针对自身
                if (val < node -> right -> val)
                {
                    node -> right = rightRotate(node -> right);
                }
                // 如果插入的节点在右子树的右边，直接左旋
                node = leftRotate(node);
            }
        }

        // 节点的高度需要更新
        updateHeight(node);
        return node;
    }

    // 计算新高度并更新
    void updateHeight(TreeNode * node)
    {
        int newHeight = max(getHeight(node -> left), getHeight(node -> right)) + 1;
        height[node] = newHeight;
    }

    // 获取节点的高度
    int getHeight(TreeNode * node)
    {
        auto it = height.find(node);
        if (it != height.end())
        {
            // 如果找到节点在map中存在 返回高度
            return it -> second;
        }
        return 0;
    }

    // 左旋
    TreeNode * leftRotate(TreeNode * node)
    {
        // 先拿到右孩子  以及右孩子的左孩子
        TreeNode * rightChild = node -> right;
        TreeNode * rightLeftChild = rightChild -> left;
        // 旋转 右孩子的左孩子是当前节点 当前节点的右孩子是右孩子的左孩子
        rightChild -> left = node;
        node -> right = rightLeftChild;

        // 节点的高度需要更新
        updateHeight(node);
        // 重新计算
        updateHeight(rightChild);
        // int nodeHeight = getHeight(node);
        // height[rightChild] = max(nodeHeight, getHeight(rightChild -> right)) + 1;

        // 返回新的根节点
        return rightChild;
    }

    // 右旋
    TreeNode * rightRotate(TreeNode * node)
    {
        // 先拿到左孩子 以及左孩子的右孩子
        TreeNode * leftChild = node -> left;
        TreeNode * leftRightChild = leftChild -> right;
        // 旋转 左孩子的右孩子是当前节点 当前节点的左孩子是左孩子的右孩子
        leftChild -> right = node;
        node -> left = leftRightChild;
        // 节点的高度需要更新
        updateHeight(node);
        updateHeight(leftChild);
        // int nodeHeight = getHeight(node);
        // height[leftChild] = max(nodeHeight, getHeight(leftChild -> left)) + 1;
        // 返回新的根节点
        return leftChild;
    }
};
